import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 数据预处理
transform = transforms.Compose([
    transforms.Resize((64, 64)),
    transforms.ToTensor(),
])

train_dataset = datasets.ImageFolder(root='data/train_data', transform=transform)
train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)

# 小型 CNN
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(3, 16, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.fc = nn.Sequential(
            nn.Flatten(),
            nn.Linear(32*16*16, 64),
            nn.ReLU(),
            nn.Linear(64, 2)  # 二分类
        )
    
    def forward(self, x):
        x = self.conv(x)
        x = self.fc(x)
        return x

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleCNN().to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练循环
epochs = 5
for epoch in range(epochs):
    model.train()
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
    print(f"Epoch {epoch+1} finished.")

# 保存模型
torch.save(model.state_dict(), "data/models/simple_cnn.pth")

# 使用模型做预测
model.eval()
sample_img, _ = train_dataset[0]  # 测试一张图片
sample_img = sample_img.unsqueeze(0).to(device)
output = model(sample_img)
pred = torch.argmax(output, 1)
# 输出详细结果
result = pred.item();# 1=wrong, 0=correct
if(result==1):
    print("预测标签: wrong 含有杂质")
elif(result==0):
    print("预测标签: correct 不含杂质")
